import pandas as pd
import numpy as np
from ortools.linear_solver import pywraplp
from collections import defaultdict


def calc_strategy_saving(strategy_type, leg1, leg2):
    """计算价差策略的保证金节省金额"""
    # 计算组合中的义务仓的单独保证金
    if strategy_type in ["bull_call", "bear_call", "straddle", "strangle"]:
        # 这些策略中的call义务仓需要保证金
        if strategy_type == "bull_call":
            alone_margin = calc_single_margin(leg2)  # leg2是义务仓
            strategy_margin = 0
        elif strategy_type == "bear_call":
            alone_margin = calc_single_margin(leg1)  # leg1是义务仓
            strategy_margin = abs(leg2["strike"] - leg1["strike"]) * leg1["multiplier"]
        elif strategy_type == "straddle":
            alone_margin = calc_single_margin(leg1) + calc_single_margin(leg2)
            strategy_margin = (
                max(calc_single_margin(leg1), calc_single_margin(leg2))
                + min(leg1["settle_price"], leg2["settle_price"]) * leg1["multiplier"]
            )
        elif strategy_type == "strangle":
            alone_margin = calc_single_margin(leg1) + calc_single_margin(leg2)
            strategy_margin = (
                max(calc_single_margin(leg1), calc_single_margin(leg2))
                + min(leg1["settle_price"], leg2["settle_price"]) * leg1["multiplier"]
            )
        return alone_margin - strategy_margin

    elif strategy_type in ["bull_put", "bear_put"]:
        # 这些策略中的put义务仓需要保证金
        if strategy_type == "bull_put":
            alone_margin = calc_single_margin(leg2)  # leg2是义务仓
            strategy_margin = abs(leg2["strike"] - leg1["strike"]) * leg1["multiplier"]
        elif strategy_type == "bear_put":
            alone_margin = calc_single_margin(leg1)  # leg1是义务仓
            strategy_margin = 0
        return alone_margin - strategy_margin


def preprocess_positions(df, underlying_price=None):
    """预处理持仓数据"""
    # 设置默认标的价
    if underlying_price is None:
        underlying_price = {
            "510050.SSE": 2.90,  # 假设50ETF价格为2.90元
            # 其他标的默认价...
        }

    # 提取需要列并重命名
    df = df[
        [
            "代码",
            "持仓",
            "到期日",
            "标的",
            "合约乘数",
            "CALLorPUT",
            "执行价",
            "现价",
            "可用买",
            "可用卖",
        ]
    ]
    df.columns = [
        "contract",
        "position",
        "expiry",
        "underlying",
        "multiplier",
        "type",
        "strike",
        "settle_price",
        "available_buy",
        "available_sell",
    ]
    df = df.astype(
        {"multiplier": float, "strike": float, "settle_price": float, "position": float}
    )

    # 添加方向列
    df["direction"] = np.where(df["position"] > 0, "long", "short")

    # 合并标的价到数据
    df["underlying_price"] = df["underlying"].map(
        lambda x: underlying_price.get(x, None)
    )

    # 过滤无效行
    return df[df["underlying_price"].notnull()]


def calc_single_margin(row):
    """计算单个义务仓的保证金"""
    if row["direction"] == "long":  # 权利仓不需要保证金
        return 0

    S = row["underlying_price"]
    K = float(row["strike"])
    M = float(row["multiplier"])
    settle = row["settle_price"]

    if row["type"] == "Call":  # 认购义务仓
        intrinsic_value = max(0, K - S)  # 虚值额
        margin = settle + max(0.12 * S - intrinsic_value, 0.07 * S)
        return margin * M

    else:  # 认沽义务仓
        intrinsic_value = max(0, S - K)  # 虚值额
        margin = settle + max(0.12 * S - intrinsic_value, 0.07 * K)
        return min(margin, K) * M


def calc_strategy_margin(strategy_type, contracts, portfolio_data):
    """计算组合策略保证金"""
    S = contracts[0]["underlying_price"]
    M = contracts[0]["multiplier"]

    if strategy_type == "bull_call":  # 认购牛市价差
        return 0

    elif strategy_type == "bear_call":  # 认购熊市价差
        long_strike = min(c["strike"] for c in contracts)
        short_strike = max(c["strike"] for c in contracts)
        return (short_strike - long_strike) * M

    elif strategy_type == "bull_put":  # 认沽牛市价差
        short_strike = max(c["strike"] for c in contracts)
        long_strike = min(c["strike"] for c in contracts)
        return (short_strike - long_strike) * M

    elif strategy_type == "bear_put":  # 认沽熊市价差
        return 0

    elif strategy_type == "straddle":  # 跨式空头
        call = next(c for c in contracts if c["type"] == "Call")
        put = next(c for c in contracts if c["type"] == "Put")

        # 单独保证金
        call_margin = calc_single_margin(call)
        put_margin = calc_single_margin(put)

        # 组合保证金
        return (
            max(call_margin, put_margin)
            + min(call["settle_price"], put["settle_price"]) * M
        )

    elif strategy_type == "strangle":  # 宽跨式空头
        call = next(
            c
            for c in contracts
            if c["type"] == "Call"
            and c["strike"] == max(c["strike"] for c in contracts)
        )
        put = next(
            c
            for c in contracts
            if c["type"] == "Put" and c["strike"] == min(c["strike"] for c in contracts)
        )

        # 单独保证金
        call_margin = calc_single_margin(call)
        put_margin = calc_single_margin(put)

        # 组合保证金
        return (
            max(call_margin, put_margin)
            + min(call["settle_price"], put["settle_price"]) * M
        )


def generate_strategies(group, group_name):
    """为分组持仓生成可能的策略组合"""
    strategies = []
    contracts = group.reset_index(drop=True)

    # 1. 提取不同类型的合约
    calls = contracts[contracts["type"] == "Call"]
    puts = contracts[contracts["type"] == "Put"]

    # 创建可用合约字典
    available_contracts = {}
    for _, row in contracts.iterrows():
        key = (row["contract"], row["direction"])
        available_contracts[key] = abs(row["position"])

    # 2. 生成垂直价差策略
    # 认购牛市价差：权利仓Call(低行权价) + 义务仓Call(高行权价)
    if len(calls) >= 2:
        long_calls = calls[calls["direction"] == "long"]
        short_calls = calls[calls["direction"] == "short"]

        for _, lc in long_calls.iterrows():
            for _, sc in short_calls.iterrows():
                if lc["strike"] < sc["strike"]:  # 权利仓行权价 < 义务仓行权价
                    contracts_needed = [
                        {
                            "contract": lc["contract"],
                            "direction": "long",
                            "quantity": 1,
                        },
                        {
                            "contract": sc["contract"],
                            "direction": "short",
                            "quantity": 1,
                        },
                    ]
                    # 检查可用性
                    if all(
                        available_contracts[(c["contract"], c["direction"])]
                        >= c["quantity"]
                        for c in contracts_needed
                    ):
                        strategy = {
                            "type": "bull_call",
                            "contracts": [
                                (lc["contract"], "long", 1),
                                (sc["contract"], "short", 1),
                            ],
                            "saving": calc_strategy_margin("bull_call", [lc, sc], None),
                        }
                        strategies.append(strategy)

    # 类似地生成其他垂直价差策略（bear_call, bull_put, bear_put）...

    # 2. 生成垂直价差策略
    # 认购牛市价差：权利仓Call(低行权价) + 义务仓Call(高行权价)
    if len(calls) >= 2:
        long_calls = calls[calls["direction"] == "long"]
        short_calls = calls[calls["direction"] == "short"]

        for _, lc in long_calls.iterrows():
            for _, sc in short_calls.iterrows():
                if lc["strike"] < sc["strike"]:  # 权利仓行权价 < 义务仓行权价
                    contracts_needed = [
                        {
                            "contract": lc["contract"],
                            "direction": "long",
                            "quantity": 1,
                        },
                        {
                            "contract": sc["contract"],
                            "direction": "short",
                            "quantity": 1,
                        },
                    ]
                    if all(
                        available_contracts[(c["contract"], c["direction"])]
                        >= c["quantity"]
                        for c in contracts_needed
                    ):
                        strategy = {
                            "type": "bull_call",
                            "contracts": [
                                (lc["contract"], "long", 1),
                                (sc["contract"], "short", 1),
                            ],
                            "saving": calc_strategy_saving("bull_call", lc, sc),
                        }
                        strategies.append(strategy)

        # 认购熊市价差：义务仓Call(低行权价) + 权利仓Call(高行权价)
        for _, sc in short_calls.iterrows():
            for _, lc in long_calls.iterrows():
                if sc["strike"] < lc["strike"]:  # 义务仓行权价 < 权利仓行权价
                    contracts_needed = [
                        {
                            "contract": sc["contract"],
                            "direction": "short",
                            "quantity": 1,
                        },
                        {
                            "contract": lc["contract"],
                            "direction": "long",
                            "quantity": 1,
                        },
                    ]
                    if all(
                        available_contracts[(c["contract"], c["direction"])]
                        >= c["quantity"]
                        for c in contracts_needed
                    ):
                        strategy = {
                            "type": "bear_call",
                            "contracts": [
                                (sc["contract"], "short", 1),
                                (lc["contract"], "long", 1),
                            ],
                            "saving": calc_strategy_saving("bear_call", sc, lc),
                        }
                        strategies.append(strategy)

    # 认沽牛市价差：权利仓Put(低行权价) + 义务仓Put(高行权价)
    if len(puts) >= 2:
        long_puts = puts[puts["direction"] == "long"]
        short_puts = puts[puts["direction"] == "short"]

        for _, lp in long_puts.iterrows():
            for _, sp in short_puts.iterrows():
                if lp["strike"] < sp["strike"]:  # 权利仓行权价 < 义务仓行权价
                    contracts_needed = [
                        {
                            "contract": lp["contract"],
                            "direction": "long",
                            "quantity": 1,
                        },
                        {
                            "contract": sp["contract"],
                            "direction": "short",
                            "quantity": 1,
                        },
                    ]
                    if all(
                        available_contracts[(c["contract"], c["direction"])]
                        >= c["quantity"]
                        for c in contracts_needed
                    ):
                        strategy = {
                            "type": "bull_put",
                            "contracts": [
                                (lp["contract"], "long", 1),
                                (sp["contract"], "short", 1),
                            ],
                            "saving": calc_strategy_saving("bull_put", lp, sp),
                        }
                        strategies.append(strategy)

        # 认沽熊市价差：义务仓Put(低行权价) + 权利仓Put(高行权价)
        for _, sp in short_puts.iterrows():
            for _, lp in long_puts.iterrows():
                if sp["strike"] < lp["strike"]:  # 义务仓行权价 < 权利仓行权价
                    contracts_needed = [
                        {
                            "contract": sp["contract"],
                            "direction": "short",
                            "quantity": 1,
                        },
                        {
                            "contract": lp["contract"],
                            "direction": "long",
                            "quantity": 1,
                        },
                    ]
                    if all(
                        available_contracts[(c["contract"], c["direction"])]
                        >= c["quantity"]
                        for c in contracts_needed
                    ):
                        strategy = {
                            "type": "bear_put",
                            "contracts": [
                                (sp["contract"], "short", 1),
                                (lp["contract"], "long", 1),
                            ],
                            "saving": calc_strategy_saving("bear_put", sp, lp),
                        }
                        strategies.append(strategy)
    # 3. 生成跨式空头策略
    # 要求：相同行权价的Call义务仓和Put义务仓
    short_calls = calls[calls["direction"] == "short"]
    short_puts = puts[puts["direction"] == "short"]

    for _, sc in short_calls.iterrows():
        for _, sp in short_puts.iterrows():
            if sc["strike"] == sp["strike"]:
                contracts_needed = [
                    {"contract": sc["contract"], "direction": "short", "quantity": 1},
                    {"contract": sp["contract"], "direction": "short", "quantity": 1},
                ]
                # 检查可用性
                if all(
                    available_contracts[(c["contract"], c["direction"])]
                    >= c["quantity"]
                    for c in contracts_needed
                ):
                    strategy = {
                        "type": "straddle",
                        "contracts": [
                            (sc["contract"], "short", 1),
                            (sp["contract"], "short", 1),
                        ],
                        "saving": calc_strategy_margin("straddle", [sc, sp], None),
                    }
                    strategies.append(strategy)

    # 4. 生成宽跨式空头策略
    # 要求：高行权价Call义务仓 + 低行权价Put义务仓
    for _, sc in short_calls.iterrows():
        for _, sp in short_puts.iterrows():
            if sc["strike"] > sp["strike"]:
                contracts_needed = [
                    {"contract": sc["contract"], "direction": "short", "quantity": 1},
                    {"contract": sp["contract"], "direction": "short", "quantity": 1},
                ]
                # 检查可用性
                if all(
                    available_contracts[(c["contract"], c["direction"])]
                    >= c["quantity"]
                    for c in contracts_needed
                ):
                    strategy = {
                        "type": "strangle",
                        "contracts": [
                            (sc["contract"], "short", 1),
                            (sp["contract"], "short", 1),
                        ],
                        "saving": calc_strategy_margin("strangle", [sc, sp], None),
                    }
                    strategies.append(strategy)

    return strategies


class OptimalMarginCalculator:
    def __init__(self, df, underlying_price=None):
        # 预处理持仓数据
        self.df = preprocess_positions(df, underlying_price)

        # 计算单独保证金
        self.df["single_margin"] = self.df.apply(calc_single_margin, axis=1)

        # 按标的和到期日分组
        grouped = self.df.groupby(["underlying", "expiry"])
        self.strategy_map = {}

        # 生成所有可能的策略
        self.all_strategies = []
        self.contract_map = {}

        # 创建合约向量
        for (underlying, expiry), group in grouped:
            group_name = f"{underlying}_{expiry}"
            strategies = generate_strategies(group, group_name)

            for strategy in strategies:
                strategy["group"] = group_name
                self.all_strategies.append(strategy)

                # 添加到合约映射
                for contract, direction, qty in strategy["contracts"]:
                    key = (contract, direction)
                    if key not in self.contract_map:
                        self.contract_map[key] = self.df[
                            self.df["contract"] == contract
                        ].iloc[0]["position"]

    def solve(self):
        solver = pywraplp.Solver.CreateSolver("SCIP")

        # 决策变量: 每种策略使用的次数
        x = {}
        for i, strategy in enumerate(self.all_strategies):
            x[i] = solver.IntVar(0, solver.infinity(), f"x_{i}")

        # 目标函数: 最大化总节省保证金
        objective = solver.Objective()
        for i, strategy in enumerate(self.all_strategies):
            objective.SetCoefficient(x[i], strategy["saving"])
        objective.SetMaximization()

        # 约束条件: 合约使用量 <= 可用持仓量
        constraints = {}
        for (contract, direction), position in self.contract_map.items():
            constraint = solver.Constraint(0, abs(position) + 0.0)
            constraints[(contract, direction)] = constraint

            # 添加策略约束
            for i, strategy in enumerate(self.all_strategies):
                for c, d, qty in strategy["contracts"]:
                    if c == contract and d == direction:
                        constraint.SetCoefficient(x[i], qty)

        # 求解
        status = solver.Solve()

        if status == pywraplp.Solver.OPTIMAL:
            solution = {}
            total_saving = 0
            used_contracts = defaultdict(float)

            for i, strategy in enumerate(self.all_strategies):
                count = x[i].solution_value()
                if count > 0:
                    solution[i] = {
                        "type": strategy["type"],
                        "count": count,
                        "saving": strategy["saving"] * count,
                        "contracts": strategy["contracts"],
                    }
                    total_saving += strategy["saving"] * count

                    # 更新已用合约数量
                    for c, d, qty in strategy["contracts"]:
                        used_contracts[(c, d)] += qty * count

            return {
                "total_saving": total_saving,
                "strategies": solution,
                "used_contracts": dict(used_contracts),
            }
        else:
            return None


# 示例使用
if __name__ == "__main__":
    # 读取持仓数据
    df = pd.read_csv("pos_local.csv", encoding="gbk")
    df = df[df["账号"] == "5235125601|5235125601(沪深期权)"]
    # 设置标的资产价格
    underlying_price = {
        "510050.SSE": 2.755,  # 假设50ETF价格为2.90元
        "510300.SSE": 3.984,
        "510500.SSE": 5.78,
        # 其他标的...
    }

    # 创建优化器
    calculator = OptimalMarginCalculator(df, underlying_price)

    # 执行优化
    result = calculator.solve()

    # 输出结果
    if result:
        print(f"最优节省保证金: {result['total_saving']:.2f}元")
        print("\n使用的策略:")
        for sid, details in result["strategies"].items():
            contracts_desc = ", ".join(
                f"{c[0]} ({c[1]}, {details['count']*c[2]}张)"
                for c in details["contracts"]
            )
            print(
                f"- 策略类型: {details['type']}, 使用次数: {details['count']}, 节省: {details['saving']:.2f}元, 合约: [{contracts_desc}]"
            )

        print("\n已使用的合约:")
        for (contract, direction), quantity in result["used_contracts"].items():
            print(f"- {contract} [{direction}] 已用数量: {quantity}张")

        # 计算总保证金对比
        total_alone_margin = -sum(
            calculator.df["single_margin"] * calculator.df["position"]
        )
        portfolio_margin = total_alone_margin - result["total_saving"]

        print("\n保证金对比:")
        print(f"单独保证金总额: {total_alone_margin:.2f}元")
        print(f"组合保证金总额: {portfolio_margin:.2f}元")
        print(
            f"节省保证金: {result['total_saving']:.2f}元 (节省率: {result['total_saving']/total_alone_margin:.1%})"
        )
    else:
        print("未找到有效解")
